import tensorflow as tf
from captcha_func import *
def crack_captcha(captcha_image):
    # 加载模型输出节点
    output = graph.get_tensor_by_name("out:0")

    with tf.Session() as sess:
        ckpt = tf.train.get_checkpoint_state('captcha_model')
        if ckpt and ckpt.model_checkpoint_path:
            checkpoint_path = ckpt.model_checkpoint_path
            saver.restore(sess, checkpoint_path)
        else:
            print('Have no mode')
            return
        predict = tf.argmax(tf.reshape(output, [-1, MAX_CAPTCHA, CHAR_SET_LEN]), 2)
        text_list = sess.run(predict, feed_dict={X: [captcha_image], keep_prob: 1})
        text = text_list[0].tolist()
        return vec2text(text)


if __name__ == '__main__':
    # 命令行参数
    tf.app.flags.DEFINE_integer('img_width', IMAGE_WIDTH, 'Image Width')
    tf.app.flags.DEFINE_integer('img_height', IMAGE_HEIGHT, 'Image Height')
    tf.app.flags.DEFINE_integer('captcha_len', MAX_CAPTCHA, 'captcha text length')
    tf.app.flags.DEFINE_integer('choice ', 3,'Choice char set，1 is only number，2 is Pure numbers plus lowercase letters，3 or other is all')
    FLAGS = tf.app.flags.FLAGS
    img_width = FLAGS.img_width
    img_height = FLAGS.img_height
    captcha_len = FLAGS.captcha_len
    choice = FLAGS.choice
    if choice == 1:
        char_set = number
    elif choice == 2:
        char_set = number + alphabet
    else:
        char_set = number + alphabet + ALPHABET
    CHAR_SET_LEN = len(char_set)
    text, image = gen_captcha_text_and_image(char_set,img_width,img_height)

    f = plt.figure()
    ax = f.add_subplot(111)
    ax.text(0.1, 0.9, text, ha='center', va='center', transform=ax.transAxes)
    plt.imshow(image)
    plt.show()
    image = convert2gray(image)
    image = image.flatten() / 255
    saver = tf.train.import_meta_graph('captcha_model/crack_capcha.model.meta', clear_devices=True)
    graph = tf.get_default_graph()
    X = graph.get_tensor_by_name("X:0")
    Y = graph.get_tensor_by_name("Y:0")
    keep_prob = graph.get_tensor_by_name("keep_prob:0")
    #调用模型预测
    predict_text = crack_captcha(image)

    print("正确: {}  预测: {}".format(text, predict_text))
